package com.vcvt.verify.utils;

import sun.misc.BASE64Decoder;
import sun.misc.BASE64Encoder;

import javax.crypto.BadPaddingException;
import javax.crypto.Cipher;
import javax.crypto.IllegalBlockSizeException;
import javax.crypto.NoSuchPaddingException;
import javax.crypto.spec.OAEPParameterSpec;
import javax.crypto.spec.PSource;
import javax.servlet.http.HttpSession;
import java.io.*;
import java.nio.charset.StandardCharsets;
import java.security.*;
import java.security.spec.MGF1ParameterSpec;
import java.security.spec.PKCS8EncodedKeySpec;
import java.security.spec.X509EncodedKeySpec;
import java.util.HashMap;
import java.util.Map;

public class RSAUtil {
    private static Cipher cipher;

    static {
        try {
            cipher = Cipher.getInstance("RSA/ECB/OAEPWITHSHA-256ANDMGF1PADDING");
        } catch (NoSuchAlgorithmException | NoSuchPaddingException e) {
            e.printStackTrace();
        }
    }

    public static Map<String, String> generateKeyPair(String filePath) {
        try (FileWriter pubfw = new FileWriter(filePath + "/publicKey.keystore");
             FileWriter prifw = new FileWriter(filePath + "/privateKey.keystore");
             BufferedWriter pubbw = new BufferedWriter(pubfw);
             BufferedWriter pribw = new BufferedWriter(prifw)) {
            KeyPairGenerator keyPairGen = KeyPairGenerator.getInstance("RSA");
            keyPairGen.initialize(2048);
            KeyPair keyPair = keyPairGen.generateKeyPair();
            PublicKey publicKey = keyPair.getPublic();
            PrivateKey privateKey = keyPair.getPrivate();
            String publicKeyString = getKeyString(publicKey);
            String privateKeyString = getKeyString(privateKey);
            pubbw.write(publicKeyString);
            pribw.write(privateKeyString);
            pubbw.flush();
            pribw.flush();
            Map<String, String> map = new HashMap<String, String>();
            map.put("publicKey", publicKeyString);
            map.put("privateKey", privateKeyString);
            return map;
        } catch (Exception e) {
            e.printStackTrace();
        }
        return null;
    }

    public static Map<String, String> generateKeyPair() {
        try {
            KeyPairGenerator keyPairGen = KeyPairGenerator.getInstance("RSA");
            keyPairGen.initialize(2048);
            KeyPair keyPair = keyPairGen.generateKeyPair();
            PublicKey publicKey = keyPair.getPublic();
            PrivateKey privateKey = keyPair.getPrivate();
            String publicKeyString = getKeyString(publicKey);
            String privateKeyString = getKeyString(privateKey);
            Map<String, String> map = new HashMap<>();
            map.put("publicKey", publicKeyString);
            map.put("privateKey", privateKeyString);
            return map;
        } catch (Exception e) {
            e.printStackTrace();
        }
        return null;
    }

    public static PublicKey getPublicKey(String key) throws Exception {
        byte[] keyBytes;
        keyBytes = (new BASE64Decoder()).decodeBuffer(key);
        X509EncodedKeySpec keySpec = new X509EncodedKeySpec(keyBytes);
        KeyFactory keyFactory = KeyFactory.getInstance("RSA");
        return keyFactory.generatePublic(keySpec);
    }

    public static PrivateKey getPrivateKey(String key) throws Exception {
        byte[] keyBytes;
        keyBytes = (new BASE64Decoder()).decodeBuffer(key);
        PKCS8EncodedKeySpec keySpec = new PKCS8EncodedKeySpec(keyBytes);
        KeyFactory keyFactory = KeyFactory.getInstance("RSA");
        return keyFactory.generatePrivate(keySpec);
    }

    public static String getKeyString(Key key) {
        byte[] keyBytes = key.getEncoded();
        return (new BASE64Encoder()).encode(keyBytes);
    }

    public static String encrypt(PublicKey publicKey, String plainText) {
        try {
            cipher.init(Cipher.ENCRYPT_MODE, publicKey);
            byte[] enBytes = cipher.doFinal(plainText.getBytes());
            return (new BASE64Encoder()).encode(enBytes);
        } catch (InvalidKeyException | IllegalBlockSizeException | BadPaddingException e) {
            e.printStackTrace();
        }
        return null;
    }

    public static String fileEncrypt(String publicKeystore, String plainText) {
        try (FileReader fr = new FileReader(publicKeystore);
             BufferedReader br = new BufferedReader(fr)) {
            StringBuilder publicKeyString = new StringBuilder();
            String str;
            while ((str = br.readLine()) != null) {
                publicKeyString.append(str);
            }
            cipher.init(Cipher.ENCRYPT_MODE, getPublicKey(publicKeyString.toString()));
            byte[] enBytes = cipher.doFinal(plainText.getBytes());
            return (new BASE64Encoder()).encode(enBytes);
        } catch (Exception e) {
            e.printStackTrace();
        }
        return null;
    }

    public static String encrypt(String publicKey, String plainText) {
        try {
            cipher.init(Cipher.ENCRYPT_MODE, getPublicKey(publicKey));
            byte[] enBytes = cipher.doFinal(plainText.getBytes(StandardCharsets.UTF_8));
            return (new BASE64Encoder()).encode(enBytes);
        } catch (Exception e) {
            e.printStackTrace();
        }
        return null;
    }

    public static String decrypt(PrivateKey privateKey, String enStr) {
        try {
            cipher.init(Cipher.DECRYPT_MODE, privateKey);
            byte[] deBytes = cipher.doFinal((new BASE64Decoder()).decodeBuffer(enStr));
            return new String(deBytes);
        } catch (InvalidKeyException | IllegalBlockSizeException | BadPaddingException | IOException e) {
            e.printStackTrace();
        }
        return null;
    }

    public static String decrypt(String privateKey, String enStr) {
        try {
            Cipher oaepFromInit = Cipher.getInstance("RSA/ECB/OAEPPadding");
            OAEPParameterSpec oaepParams = new OAEPParameterSpec("SHA-256", "MGF1",
                    new MGF1ParameterSpec("SHA-1"), PSource.PSpecified.DEFAULT);
            oaepFromInit.init(Cipher.DECRYPT_MODE, getPrivateKey(privateKey), oaepParams);
            byte[] deBytes = oaepFromInit.doFinal((new BASE64Decoder()).decodeBuffer(enStr));
            return new String(deBytes);
        } catch (Exception e) {
            e.printStackTrace();
        }
        return null;
    }

    public static String fileDecrypt(String privateKeystore, String enStr) {
        try (FileReader fr = new FileReader(privateKeystore);
             BufferedReader br = new BufferedReader(fr)) {
            String privateKeyString = "";
            String str;
            while ((str = br.readLine()) != null) {
                privateKeyString += str;
            }
            cipher.init(Cipher.DECRYPT_MODE, getPrivateKey(privateKeyString));
            byte[] deBytes = cipher.doFinal((new BASE64Decoder()).decodeBuffer(enStr));
            return new String(deBytes);
        } catch (Exception e) {
            e.printStackTrace();
        }
        return null;
    }


/*    *//**
     * 样例
     * @param args
     * @throws Exception
     */
    public static void main(String[] args) throws Exception {
        String password = "123456";
        Map<String, String> maps = generateKeyPair();
        PublicKey pubkey = getPublicKey(maps.get("publicKey"));
        PrivateKey privkey = getPrivateKey(maps.get("privateKey"));
// --- encrypt given algorithm string
        Cipher oaepFromAlgo = Cipher.getInstance("RSA/ECB/OAEPWITHSHA-256ANDMGF1PADDING");
        oaepFromAlgo.init(Cipher.ENCRYPT_MODE, pubkey);
        byte[] ct = oaepFromAlgo.doFinal(password.getBytes(StandardCharsets.UTF_8));

// --- decrypt given OAEPParameterSpec
        Cipher oaepFromInit = Cipher.getInstance("RSA/ECB/OAEPPadding");
        OAEPParameterSpec oaepParams = new OAEPParameterSpec("SHA-256", "MGF1", new MGF1ParameterSpec("SHA-1"), PSource.PSpecified.DEFAULT);
        oaepFromInit.init(Cipher.DECRYPT_MODE, privkey, oaepParams);
        byte[] pt = oaepFromInit.doFinal(ct);
        System.out.println(new String(pt, StandardCharsets.UTF_8));


        /**
         * 通过generateKeyPair获取公钥和私钥返回给前端，前端代码，先用npm install node-forge安装
         * this.publickey = '-----BEGIN PUBLIC KEY-----\n' + this.publickey + '\n-----END PUBLIC KEY-----'
         *             var forge = require('node-forge');
         *             const publicK = forge.pki.publicKeyFromPem(this.publickey)
         *             var encrypted =publicK.encrypt(pass, 'RSA-OAEP', {
         *               md: forge.md.sha256.create(),
         *               mgf1: {
         *                 md: forge.md.sha1.create()
         *               }
         *             });
         *             this.loginForm.passwordEncryption = forge.util.encode64(encrypted)
         *             this.loginForm.privateKey = this.privateKey
         *  将passwordEncryption和 privateKey传给后端进行解密
         *
         */
    }
}
